from typing import List, Dict, Optional
from mcp.server.fastmcp import FastMCP
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse

import uvicorn
import json
import asyncio
import uuid
import time

# 初始化 MCP 实例
mcp = FastMCP()

# 示例工具函数
@mcp.tool()
async def hello_world(name: str = "World") -> Dict:
    """一个简单的示例工具，返回问候语"""
    return {"message": f"Hello, {name}!"}

@mcp.tool()
async def add_numbers(a: float, b: float) -> Dict:
    """将两个数字相加"""
    return {"result": a + b}

@mcp.tool()
async def list_items(items: List[str]) -> Dict:
    """返回输入列表的信息"""
    return {
        "count": len(items),
        "items": items
    }

@mcp.tool()
async def get_weather(city: str) -> Dict:
    """获取指定城市的天气信息"""
    # 模拟数据，实际应用中可以调用天气 API
    weather_data = {
        "北京": {"temperature": "25°C", "condition": "晴朗", "humidity": "40%"},
        "上海": {"temperature": "28°C", "condition": "多云", "humidity": "65%"},
        "广州": {"temperature": "30°C", "condition": "阵雨", "humidity": "80%"},
        "深圳": {"temperature": "29°C", "condition": "多云", "humidity": "75%"},
    }
    
    if city in weather_data:
        return {
            "city": city,
            "weather": weather_data[city],
            "forecast": "未来几天天气良好"
        }
    else:
        return {
            "city": city,
            "error": "没有找到该城市的天气信息"
        }

# 创建应用
app = FastAPI(
    title="MCP SSE Server Example",
    description="A simple MCP server with SSE transport"
)

# 添加 CORS 中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

async def get_tools_json():
    """获取工具列表的JSON格式"""
    tools = await mcp.list_tools()
    return [
        {
            "name": tool.name,
            "description": tool.description,
            "inputSchema": tool.inputSchema
        }
        for tool in tools
    ]

# 为了简化，直接使用这些路由
@app.get("/sse")
async def sse_endpoint(request: Request):
    """SSE 端点，仅发送最基本的消息"""
    print(f"Client connected to SSE from {request.client}")
    
    async def event_generator():
        # 1. 初始化响应
        yield "data: {\"jsonrpc\":\"2.0\",\"id\":\"1\",\"result\":{\"type\":\"initialize\"}}\n\n"
        await asyncio.sleep(0.5)
        
        # 2. 发送工具列表
        tools_json = await get_tools_json()
        tools_msg = json.dumps({
            "jsonrpc": "2.0", 
            "id": "2", 
            "result": {"tools": tools_json}
        })
        yield f"data: {tools_msg}\n\n"
        print(f"Sent tools list to client: {len(tools_json)} tools")
        
        # 3. 持续发送心跳
        count = 0
        while True:
            if await request.is_disconnected():
                print("Client disconnected from SSE")
                break
                
            # 每5秒发送一次心跳
            count += 1
            heartbeat = json.dumps({
                "jsonrpc": "2.0",
                "method": "heartbeat",
                "params": {"count": count, "time": time.time()}
            })
            print(f"Sending heartbeat #{count}")
            yield f"data: {heartbeat}\n\n"
            await asyncio.sleep(5)
    
    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream"
    )

@app.post("/call-tool")
async def call_tool(request: Request):
    """简化的工具调用端点"""
    try:
        data = await request.json()
        print(f"Received call_tool request: {data}")
        
        # 尝试提取工具名称和参数
        if "method" in data and data["method"] == "call_tool" and "params" in data:
            # JSON-RPC 2.0 格式
            params = data["params"]
            tool_name = params.get("name")
            tool_args = params.get("parameters", {})
            req_id = data.get("id", "1")
        else:
            # 简单格式
            tool_name = data.get("name")
            tool_args = data.get("parameters", {})
            req_id = data.get("id", "1")
        
        if not tool_name:
            return {"jsonrpc": "2.0", "id": req_id, "error": {"message": "Tool name is required"}}
        
        print(f"Calling tool '{tool_name}' with args: {tool_args}")
        result = await mcp.call_tool(tool_name, tool_args)
        print(f"Tool result: {result}")
        
        return {
            "jsonrpc": "2.0",
            "id": req_id,
            "result": result
        }
    except Exception as e:
        print(f"Error in call_tool: {str(e)}")
        return {
            "jsonrpc": "2.0",
            "id": "error",
            "error": {"message": str(e)}
        }

@app.get("/health")
async def health_check():
    """健康检查端点"""
    return {"status": "healthy"}

if __name__ == "__main__":
    print("Starting MCP SSE Server...")
    print("Server running at http://localhost:8000")
    print("SSE endpoint at http://localhost:8000/sse")
    uvicorn.run(app, host="0.0.0.0", port=8000) 